package socks5

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"strconv"
	"sync"
)

var bufs = sync.Pool{New: func() any { return make([]byte, 255) }}

type Auth struct {
	User string
	Pwd  string
}

type Proxy struct {
	addr string
	auth *Auth
}

func New(addr string, auth *Auth) *Proxy {
	return &Proxy{addr: addr, auth: auth}
}

func (p *Proxy) Connect(addr string) (c net.Conn, err error) {
	return p.Dial("tcp", addr)
}

func (p *Proxy) Dial(network, addr string) (c net.Conn, err error) {
	return p.DialContext(context.Background(), network, addr)
}

func (p *Proxy) DialContext(ctx context.Context, network, addr string) (c net.Conn, err error) {
	host, sp, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	port, err := strconv.Atoi(sp)
	if err != nil {
		return nil, err
	}

	if len(host) > 255 || port < 1 || port > 65535 {
		return nil, fmt.Errorf("invalid address")
	}
	if p.auth != nil && (len(p.auth.User) > 255 || len(p.auth.Pwd) > 255) {
		return nil, fmt.Errorf("invalid auth")
	}

	var nd net.Dialer
	conn, err := nd.DialContext(ctx, network, p.addr)
	if err != nil {
		return nil, err
	}

	var ret net.Conn
	defer func() {
		if ret == nil {
			conn.Close()
		}
	}()
	buf := bufs.Get().([]byte)
	defer bufs.Put(buf)
	read := func(size byte) ([]byte, error) {
		b := buf[:size]
		if _, err := io.ReadFull(conn, b); err != nil {
			return nil, err
		}
		return b, nil
	}

	ver := byte(0x05)  // 只处理Socks5协议
	auth := byte(0x00) // 不需要认证
	if p.auth != nil {
		auth = 0x02 // 账号密码认证
	}
	b := []byte{
		ver,
		0x01, // METHODS_COUNT
		auth,
	}
	if _, err := conn.Write(b); err != nil {
		return nil, err
	}

	b, err = read(2) // 服务端返回选择的认证方法
	if err != nil {
		return nil, err
	}
	if b[0] != ver || b[1] != auth {
		return nil, fmt.Errorf("no acceptable methods")
	}

	if auth == 0x02 {
		usr := []byte(p.auth.User)
		pwd := []byte(p.auth.Pwd)
		b = make([]byte, len(usr)+len(pwd)+3)
		i := 0
		b[i] = 0x01 // 认证子协商版本
		i += 1
		b[i] = byte(len(usr)) // USERNAME_LENGTH
		i += 1
		copy(b[i:], usr) // USERNAME
		i += len(usr)
		b[i] = byte(len(pwd)) // PASSWORD_LENGTH
		i += 1
		copy(b[i:], pwd) // PASSWORD
		if _, err := conn.Write(b); err != nil {
			return nil, err
		}

		b, err = read(2) // 服务端响应账号密码认证结果
		if err != nil {
			return nil, err
		}
		if b[0] != 0x01 || b[1] != 0x00 { // 0x00 认证成功
			return nil, fmt.Errorf("auth failed")
		}
	}

	ht, hb := parseHost(host)
	b = make([]byte, len(hb)+6)
	i := 0
	b[i] = ver
	i += 1
	b[i] = 0x01 // CONNECT 连接上游服务器
	i += 2
	// b[i] = 0x00 //RSV 保留字段
	b[i] = ht // 地址类型
	i += 1
	copy(b[i:], hb)
	i += len(hb)
	binary.BigEndian.PutUint16(b[i:], uint16(port))
	if _, err := conn.Write(b); err != nil {
		return nil, err
	}

	b, err = read(4) // 响应连接头 VERSION RESPONSE RSV ADDRESS_TYPE
	if err != nil {
		return nil, err
	}
	if b[0] != ver || b[1] != 0x00 { // 0x00 代理服务器连接目标服务器成功
		return nil, fmt.Errorf("connect host failed, code %d", b[1])
	}
	switch b[3] { // ADDRESS_TYPE
	case 0x01: // IPv4
		if _, err := read(4); err != nil {
			return nil, err
		}
	case 0x03: // 域名
		size, err := read(1)
		if err != nil {
			return nil, err
		}
		if _, err := read(size[0]); err != nil {
			return nil, err
		}
	case 0x04: // IPv6
		if _, err := read(16); err != nil {
			return nil, err
		}
	}
	if _, err := read(2); err != nil { // port
		return nil, err
	}

	if ctx.Err() != nil {
		return nil, err
	}

	ret = conn
	return ret, nil
}

func parseHost(s string) (byte, []byte) {
	if ip := net.ParseIP(s); ip != nil {
		if ip4 := ip.To4(); ip.Equal(ip4) {
			return 0x01, ip4 // IPv4
		} else {
			return 0x04, ip.To16() // IPv6
		}
	}
	b := []byte(s)
	return 0x03, append([]byte{byte(len(b))}, b...) // Domain
}

//

type Conn struct {
	conn net.Conn
	auth func(usr, pwd []byte) bool
}

func NewConn(conn net.Conn, auth func(usr, pwd []byte) bool) *Conn {
	return &Conn{conn: conn, auth: auth}
}

func copyBytes(b []byte) []byte {
	ret := make([]byte, len(b))
	copy(ret, b)
	return ret
}

func (c *Conn) ReadConnect() (string, error) {
	buf := bufs.Get().([]byte)
	defer bufs.Put(buf)
	read := func(size byte) ([]byte, error) {
		b := buf[:size]
		if _, err := io.ReadFull(c.conn, b); err != nil {
			return nil, err
		}
		return b, nil
	}

	// ctx, cancel := context.WithTimeout(app.CTX, time.Second*3)
	// defer cancel()
	// go func() {
	// 	<-ctx.Done()
	// 	if ctx.Err() == context.DeadlineExceeded {
	// 		conn.Close()
	// 	}
	// }()

	makeErr := func(s string) (string, error) {
		return "", fmt.Errorf("ReadHeader %s: %s", c.conn.RemoteAddr(), s)
	}

	b, err := read(2) // VER NMETHODS
	if err != nil {
		return makeErr(err.Error())
	}
	if b[0] != 0x05 { // 只处理Socks5协议
		return makeErr("协议版本错误: " + strconv.Itoa(int(b[0])))
	}
	b, err = read(b[1]) // METHODS
	if err != nil {
		return makeErr(err.Error())
	}
	// log.Debug(b)

	acceptMethod := byte(0x00) // 无验证
	if c.auth != nil {
		acceptMethod = byte(0x02) // 密码验证
	}
	if !bytes.Contains(b, []byte{acceptMethod}) {
		c.conn.Write([]byte{0x05, 0xFF}) // 回应客户端：无可接受的方法
		return makeErr("验证方式错误: " + fmt.Sprint(b))
	}
	c.conn.Write([]byte{0x05, acceptMethod}) // 回应客户端：Socks服务端不需要验证方式

	if c.auth != nil { // 密码验证
		b, err = read(1) // VERSION
		if err != nil {
			return makeErr(err.Error())
		}
		if b[0] != 0x01 {
			c.conn.Write([]byte{0x01, 0x01}) // 回应客户端：认证失败
			return makeErr("密码验证版本错误: " + strconv.Itoa(int(b[0])))
		}
		b, err = read(1)
		if err != nil {
			return makeErr(err.Error())
		}
		b, err = read(b[0]) // username
		if err != nil {
			return makeErr(err.Error())
		}
		usr := copyBytes(b)
		b, err = read(1)
		if err != nil {
			return makeErr(err.Error())
		}
		b, err = read(b[0]) // password
		if err != nil {
			return makeErr(err.Error())
		}
		pwd := copyBytes(b)
		if c.auth(usr, pwd) {
			c.conn.Write([]byte{0x01, 0x00}) // 回应客户端：认证成功
		} else {
			c.conn.Write([]byte{0x01, 0x01}) // 回应客户端：认证失败
			return makeErr("密码验证失败")
		}
	}

	b, err = read(4) // 命令请求头 VERSION COMMAND RSV ADDRESS_TYPE
	if err != nil {
		return makeErr(err.Error())
	}
	if b[0] != 0x05 || b[1] != 0x01 { // 只处理Socks5协议 CONNECT命令
		return makeErr("指令错误: " + fmt.Sprint(b))
	}

	var host string
	switch b[3] { // ADDRESS_TYPE
	case 0x01: // IP V4
		b, err := read(4)
		if err != nil {
			return makeErr(err.Error())
		}
		host = net.IP(b).String()
	case 0x03: // 域名
		size, err := read(1)
		if err != nil {
			return makeErr(err.Error())
		}
		b, err := read(size[0])
		if err != nil {
			return makeErr(err.Error())
		}
		host = string(b)
	case 0x04: // IP V6
		b, err := read(16)
		if err != nil {
			return makeErr(err.Error())
		}
		host = net.IP(b).String()
	default:
		return makeErr("地址类型错误")
	}
	b, err = read(2)
	if err != nil {
		return makeErr(err.Error())
	}
	port := int(binary.BigEndian.Uint16(b))

	return net.JoinHostPort(host, strconv.Itoa(port)), nil
}

func (c *Conn) WriteConnectStatus(success bool) error {
	b := []byte{0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00} // 响应客户端连接成功
	if !success {
		b[1] = 0x01 // general SOCKS server failure
	}
	_, err := c.conn.Write(b) // 响应客户端连接状态
	return err
}
